{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e22d32fa",
   "metadata": {},
   "source": [
    "# Optimization with Linopy - Migrate Extra Functionalities\n",
    "\n",
    "The extra funcionalities for the native pypsa optimization code are mostly using the function of the `pypsa.linopt` model. Here we show how you can recycle large parts of your code by using the compatibility functions from the `pypsa.optimization.compat` module.\n",
    "\n",
    "These are \n",
    "\n",
    "* `define_variables`\n",
    "* `define_constraints`\n",
    "* `get_var`\n",
    "* `linexpr`\n",
    "\n",
    "You might want to use them if you have `extra_functionalities` written for the native optimization code. However, expect some hick-ups, as some operations might behave differently. \n",
    "\n",
    "Let's import pypsa and the compat functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b59d1a94",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pypsa\n",
    "from pypsa.optimization.compat import get_var, define_constraints, linexpr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50959aff",
   "metadata": {},
   "source": [
    "We load the same network from the previous section into memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d24d81a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = pypsa.examples.ac_dc_meshed(from_master=True)\n",
    "\n",
    "n.generators.loc[n.generators.carrier == \"gas\", \"p_nom_extendable\"] = False\n",
    "n.generators.loc[n.generators.carrier == \"gas\", \"ramp_limit_down\"] = 0.2\n",
    "n.generators.loc[n.generators.carrier == \"gas\", \"ramp_limit_up\"] = 0.2\n",
    "\n",
    "n.add(\n",
    "    \"StorageUnit\",\n",
    "    \"su\",\n",
    "    bus=\"Manchester\",\n",
    "    marginal_cost=10,\n",
    "    inflow=50,\n",
    "    p_nom_extendable=True,\n",
    "    capital_cost=10,\n",
    "    p_nom=2000,\n",
    "    efficiency_dispatch=0.5,\n",
    "    cyclic_state_of_charge=True,\n",
    "    state_of_charge_initial=1000,\n",
    ")\n",
    "\n",
    "n.add(\n",
    "    \"StorageUnit\",\n",
    "    \"su2\",\n",
    "    bus=\"Manchester\",\n",
    "    marginal_cost=10,\n",
    "    p_nom_extendable=True,\n",
    "    capital_cost=50,\n",
    "    p_nom=2000,\n",
    "    efficiency_dispatch=0.5,\n",
    "    carrier=\"gas\",\n",
    "    cyclic_state_of_charge=False,\n",
    "    state_of_charge_initial=1000,\n",
    ")\n",
    "\n",
    "n.storage_units_t.state_of_charge_set.loc[n.snapshots[7], \"su\"] = 100\n",
    "\n",
    "n.add(\"Bus\", \"storebus\", carrier=\"hydro\", x=-5, y=55)\n",
    "n.madd(\n",
    "    \"Link\",\n",
    "    [\"battery_power\", \"battery_discharge\"],\n",
    "    \"\",\n",
    "    bus0=[\"Manchester\", \"storebus\"],\n",
    "    bus1=[\"storebus\", \"Manchester\"],\n",
    "    p_nom=100,\n",
    "    efficiency=0.9,\n",
    "    p_nom_extendable=True,\n",
    "    p_nom_max=1000,\n",
    ")\n",
    "n.madd(\n",
    "    \"Store\",\n",
    "    [\"store\"],\n",
    "    bus=\"storebus\",\n",
    "    e_nom=2000,\n",
    "    e_nom_extendable=True,\n",
    "    marginal_cost=10,\n",
    "    capital_cost=10,\n",
    "    e_nom_max=5000,\n",
    "    e_initial=100,\n",
    "    e_cyclic=True,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e47dcffc",
   "metadata": {},
   "source": [
    "And define the extra functionalities as we defined them for the native code in [here](https://pypsa.readthedocs.io/en/latest/examples/lopf_with_pyomo_False.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1dd81f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def minimal_state_of_charge(n, snapshots):\n",
    "    vars_soc = get_var(n, \"StorageUnit\", \"state_of_charge\")\n",
    "    lhs = linexpr((1, vars_soc))\n",
    "    define_constraints(n, lhs, \">\", 50, \"StorageUnit\", \"soc_lower_bound\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c528797",
   "metadata": {},
   "source": [
    "With the compat functions, this will work as expected. Let's go on to the next one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad6b8896",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_link_cap_ratio(n, snapshots):\n",
    "    vars_link = get_var(n, \"Link\", \"p_nom\")\n",
    "    eff = n.links.at[\"battery_power\", \"efficiency\"]\n",
    "    lhs = linexpr(\n",
    "        (1, vars_link.loc[\"battery_power\"]), (-eff, vars_link.loc[\"battery_discharge\"])\n",
    "    )\n",
    "    define_constraints(n, lhs, \"=\", 0, \"battery_discharge\", attr=\"fixratio\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07526625",
   "metadata": {},
   "source": [
    "This function as well should not make any problems. Let's go on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a09f337",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_bus_production(n, snapshots):\n",
    "    total_demand = n.loads_t.p_set.sum().sum()\n",
    "    prod_per_bus = (\n",
    "        linexpr((1, get_var(n, \"Generator\", \"p\")))\n",
    "        .groupby(n.generators.bus, axis=1)\n",
    "        .apply(join_exprs)\n",
    "    )\n",
    "    define_constraints(\n",
    "        n, prod_per_bus, \">=\", total_demand / 5, \"Bus\", \"production_share\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0737b26e",
   "metadata": {},
   "source": [
    "Here, we come into difficult terrain. The groupby function won't work since the `linexpr` functions returns some sort of a xarray object (a `LinearExpression` object, derived from `xarray.Dataset`). \n",
    "\n",
    "Instead, we have to rewrite parts: \n",
    "\n",
    "* remove the `axis` argument\n",
    "* use a xaray DataArray as argument\n",
    "* explicitly sum over all snapshots afterwards. This has nothing to do with groupby but with the fact that we want to limit to total production.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e501961",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_bus_production(n, snapshots):\n",
    "    total_demand = n.loads_t.p_set.sum().sum()\n",
    "    prod_per_bus = (\n",
    "        linexpr((1, get_var(n, \"Generator\", \"p\")))\n",
    "        .groupby(n.generators.bus.to_xarray())\n",
    "        .sum()\n",
    "        .sum()\n",
    "    )\n",
    "    define_constraints(\n",
    "        n, prod_per_bus, \">=\", total_demand / 5, \"Bus\", \"production_share\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d72ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extra_functionalities(n, snapshots):\n",
    "    minimal_state_of_charge(n, snapshots)\n",
    "    fix_link_cap_ratio(n, snapshots)\n",
    "    fix_bus_production(n, snapshots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c326ac6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "n.optimize(\n",
    "    extra_functionality=extra_functionalities,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "",
   "language": "python",
   "name": ""
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
